from typing import List, Dict
from tqdm import tqdm

from hj_dataset_devkit import Scene, Obstacle

from obj_autolabeling import ObjLabeler, build_component
from components.obj_detector import ObjDetector
from components.obj_tracker import ObjTracker, FrameInfo, TrackObject

__all__ = ['DetTrackRefineLabeler']

class DetTrackRefineLabeler(ObjLabeler):
    def __init__(self, detector: dict, tracker: dict, refiner: dict) -> None:
        super().__init__()
        self.__detector: ObjDetector = build_component(detector['name'],
                                                       detector['kwargs'])
        self.__tracker: ObjTracker = build_component(tracker['name'],
                                                     tracker['kwargs'])

    def process(self, scene: Scene) -> Dict[int, List[Obstacle]]:
        frames_info = self.detect(scene)
        tracks = self.track(frames_info)
        # return TODO
        return {
            stamp: [ Obstacle(i.id, i.category, i.bbox.center.to_list(),
                              i.bbox.size.to_list(), i.bbox.rotation.to_list(),
                              0, 0) for i in tracked_obj ]
            for stamp, tracked_obj in tracks.items()
        }

    def detect(self, scene: Scene, batch_size=8) -> List[FrameInfo]:
        result = []
        if batch_size <= 1:
            for i, frame in enumerate(tqdm(scene.frames, 'detect frames')):
                result.append(FrameInfo(frame.get_timestamp(), frame.ego_pose,
                                        i, self.__detector.process_frame(frame)))
        else:
            frames_det = self.__detector.process_frames_by_batch(scene.frames, batch_size)
            for i, frame in enumerate(scene.frames):
                stamp = frame.get_timestamp()
                if stamp in frames_det.keys():
                    result.append(FrameInfo(stamp, frame.ego_pose, i, frames_det[stamp]))
        return result

    def track(self, frames_info: List[FrameInfo]) -> Dict[int, List[TrackObject]]:
        return self.__tracker.process(frames_info)
